package cn.whale.handler;

import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;

@Slf4j
public class StreamingChatResponseHandlerImpl implements StreamingChatResponseHandler {

    private SseEmitter emitter;

    public StreamingChatResponseHandlerImpl(SseEmitter emitter){
        this.emitter = emitter;
        ServletRequestAttributes requestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletResponse response = requestAttributes.getResponse();
        response.setContentType("text/event-stream;charset=UTF-8");
        response.setCharacterEncoding("UTF-8");
    }

    @Override
    public void onPartialResponse(String partialResponse) {
        try {
            log.info("获取到内容 {}",partialResponse);
            emitter.send(partialResponse);
        } catch (IOException e) {
            log.error("发生异常 {}",e.getMessage());
            emitter.complete();
        }
    }


    @Override
    public void onCompleteResponse(ChatResponse completeResponse) {
        log.debug("结束对话");
    }

    @Override
    public void onError(Throwable e) {
        log.error("发生异常 {}",e.getMessage());
    }
}
